package org.honeynet.hbbackend.submithttp;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.math.BigInteger;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Collection;
import java.util.EnumSet;
import java.util.HashMap;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import javax.ejb.ApplicationException;
import javax.ejb.EJBException;
import javax.ejb.SessionContext;
import javax.ejb.Stateless;
import javax.ejb.TransactionAttribute;
import javax.ejb.TransactionAttributeType;
import javax.jms.BytesMessage;
import javax.jms.ConnectionFactory;
import javax.jms.Destination;
import javax.jms.JMSException;
import javax.jms.Message;
import javax.jms.MessageProducer;
import javax.jms.Session;
import javax.jms.TextMessage;
import javax.resource.ResourceException;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.Part;
import javax.sql.DataSource;

import org.codehaus.jackson.JsonGenerationException;
import org.codehaus.jackson.map.JsonMappingException;
import org.codehaus.jackson.map.ObjectMapper;
import org.honeynet.hbbackend.common.storage.Storage;
import org.honeynet.hbbackend.common.types.common.AttackInfoType;
import org.honeynet.hbbackend.common.types.submithttp.json.AttackRecord;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.xadisk.bridge.proxies.interfaces.XAFileOutputStream;
import org.xadisk.connector.outbound.XADiskConnection;
import org.xadisk.connector.outbound.XADiskConnectionFactory;
import org.xadisk.filesystem.exceptions.ClosedStreamException;
import org.xadisk.filesystem.exceptions.FileAlreadyExistsException;
import org.xadisk.filesystem.exceptions.FileNotExistsException;
import org.xadisk.filesystem.exceptions.FileUnderUseException;
import org.xadisk.filesystem.exceptions.InsufficientPermissionOnFileException;
import org.xadisk.filesystem.exceptions.LockingFailedException;
import org.xadisk.filesystem.exceptions.NoTransactionAssociatedException;


//TODO caching

@Stateless
public class SubmitHttpHandler {
	
	private final static int BUF_SIZE = 8192;
	
	private final static int MAX_FORM_FIELD_SIZE = 128;
	private final static String BINARY_PART_NAME = "data";
	private final static String BINARY_CONTENT_TYPE = "application/octet-stream";
	
	private static Logger log = LoggerFactory.getLogger(SubmitHttpHandler.class);
	
	@ApplicationException(rollback=true)
	public static class SubmissionValidationException extends Exception {
		private static final long serialVersionUID = 1L;

		public SubmissionValidationException(String message) {
			super(message);
		}
	}
	
	@Resource(mappedName="jdbc/hbbackend")
	private DataSource ds;
	
	@Resource(mappedName="jms/new_attack")
	private Destination jms_new_attack;
	
	@Resource(mappedName="jms/new_binary")
	private Destination jms_new_binary;
	
	@Resource(mappedName="jms/ConnectionFactory")
	private ConnectionFactory jmsConnectionFactory;
	
	@Resource(mappedName="xadisk/ConnectionFactory")
	private XADiskConnectionFactory xaDiskConnectionFactory;
	
	@Resource(name="storageDir")
	private String storageDir;
	
	@Resource
	private SessionContext sctx;
	
	
	@PostConstruct
	public void postConstruct() {
		log.debug("storageDir = {}", storageDir);
	}
	
	
	@TransactionAttribute(TransactionAttributeType.REQUIRES_NEW)
	public SubmitHttpReturnType handle(HttpServletRequest request) throws SubmissionValidationException {
		log.trace("handle() called");
		
		String emsg;
		
		try {
			Part binaryPart = request.getPart(BINARY_PART_NAME);
			
			if (binaryPart != null) {
				return handleSecondPhase(request);
			}
			else {
				return handleFirstPhase(request);
			}
		} catch (IOException e) {
			emsg = "got IOException while trying read part: " + BINARY_PART_NAME;
			log.error(emsg, e);
			throw new EJBException(emsg, e);
		} catch (ServletException e) {
			emsg = "got ServletException while trying to read part: " + BINARY_PART_NAME;
			log.error(emsg, e);
			throw new EJBException(emsg, e);
		} 
	}
	
	
	private SubmitHttpReturnType handleFirstPhase(HttpServletRequest request) throws SubmissionValidationException {
		log.info("submission phase 1");
		
		String emsg;
		
		Collection<Part> parts = null;
		try {
			parts = request.getParts();
		} catch (IOException e) {
			emsg = "got IOException while trying to read parts"; 
			log.error(emsg, e);
			throw new EJBException(emsg, e);
		} catch (ServletException e) {
			emsg = "got ServletException while trying to read parts"; 
			log.error(emsg, e);
			throw new EJBException(emsg, e);
		}
		
		HashMap<String, String> form = new HashMap<String, String>();
		EnumSet<FormFields> fieldsToCheck = EnumSet.allOf(FormFields.class);

		for (Part part : parts) {
			String partName = part.getName();
			
			FormFields field;
			try {
				field = FormFields.valueOf(partName);
			} catch (IllegalArgumentException e) {
				log.warn("unknown field: {}", partName);
				continue;
			}
			
			// valid field, but already removed from set 
			if (fieldsToCheck.contains(field) == false) {
				log.warn("duplicate field: {}", partName);
				continue;
			}
			
			// check field size
			long partSize = part.getSize();
			if (partSize > MAX_FORM_FIELD_SIZE) {
				emsg = "field too large: " + partName;
				log.error(emsg);
				throw new SubmissionValidationException(emsg);
			}
			
			// read field
			byte[] buf = new byte[(int) partSize];
			InputStream is;
			try {
				is = part.getInputStream();
				is.read(buf);
				is.close();
			} catch (IOException e) {
				emsg = "got IOException while reading part: " + partName;
				log.error(emsg);
				throw new EJBException(emsg, e);
			}
			String partContent = new String(buf);
			log.trace("read field: {} = {}", partName, partContent);
			
			// validate
			if (field.validate(partContent) == false) {
				emsg = "field does not validate: " + partName;
				log.error(emsg);
				throw new SubmissionValidationException(emsg);
			}
			
			// store
			form.put(partName, partContent);
			
			// finally, remove from checklist
			fieldsToCheck.remove(field);
		}
		
		if (fieldsToCheck.isEmpty() == false) {
			StringBuilder missingFields = new StringBuilder();
			String separator = ", ";
			for (FormFields field : fieldsToCheck) {
				missingFields.append(field + separator);
			}
			missingFields.setLength(missingFields.length() - separator.length());
			emsg = "missing field: " + missingFields;
			log.error(emsg);
			throw new SubmissionValidationException(emsg);
		}

		// convert ipn to string
		form.put(FormFields.source_host.name(), ipnToString(form.get(FormFields.source_host.name())));
		form.put(FormFields.target_host.name(), ipnToString(form.get(FormFields.target_host.name())));
		
		log.trace("parsed all fields: {}", form);
		
		
		Connection db = null;
		javax.jms.Connection mq = null;
		Session sess = null;
		MessageProducer prod = null;
		
		PreparedStatement pStmt;
		ResultSet queryRes;
		
		
		try {
			String md5 = form.get(FormFields.md5.name());
			log.debug("storing attack in the db (md5 = {})", md5);
			
			db = ds.getConnection();

			// insert record for binary
			pStmt = db.prepareStatement("select * from safe_insert_binary(?,?,?,?,?)");
			pStmt.setString(1, md5);
			pStmt.setString(2, form.get(FormFields.sha512.name()));
			pStmt.setNull(3, java.sql.Types.VARCHAR);
			pStmt.setNull(4, java.sql.Types.INTEGER);
			pStmt.setBoolean(5, false);
			queryRes = pStmt.executeQuery();
			
			queryRes.next();
			long binaryId = queryRes.getLong(1); // ret_id
			boolean binaryStored = queryRes.getBoolean(3); // ret_stored
			if (queryRes.getBoolean(2)) { // ret_inserted
				log.debug("inserted binary (id = {}, stored = {}, md5 = {})", new Object[] { binaryId, binaryStored, md5 });
			}
			else {
				log.debug("found existing binary (id = {}, stored = {}, md5 = {})", new Object[] { binaryId, binaryStored, md5 });
			}
			queryRes.close();
			pStmt.close();
			
			// if binary is not in store yet, register (again) that the upload is pending
			if (binaryStored == false) {
				pStmt = db.prepareStatement("insert into binaries_pending (binary_id, md5) values (?,?)");
				pStmt.setLong(1, binaryId);
				pStmt.setString(2, md5);
				pStmt.executeUpdate();
				pStmt.close();
				log.debug("registered binary as pending (id = {}, md5 = {})", binaryId, md5);
			}
			
			 // insert attack
			pStmt = db.prepareStatement(
					"insert into attacks_submithttp (binary_id, sensor_ip, source_ip, source_port, target_ip, target_port, url, trig)" +
					" values (?,inet(?),inet(?),?,inet(?),?,?,?)" +
					" returning id"
					);
			pStmt.setLong(1, binaryId);
			pStmt.setString(2, "127.0.0.1"); //TODO
			pStmt.setString(3, form.get(FormFields.source_host.name()));
			pStmt.setInt(4, Integer.parseInt(form.get(FormFields.source_port.name())));
			pStmt.setString(5, form.get(FormFields.target_host.name()));
			pStmt.setInt(6, Integer.parseInt(form.get(FormFields.target_port.name())));
			pStmt.setString(7, form.get(FormFields.url.name()));
			pStmt.setString(8, form.get(FormFields.trigger.name()));
			queryRes = pStmt.executeQuery();
			queryRes.next();
			long attackId = queryRes.getLong(1);
			queryRes.close();
			pStmt.close();
			log.debug("inserted attack (id = {}, binary_id = {}, md5 = {})", new Object[] { attackId, binaryId, md5 });
			
			mq = jmsConnectionFactory.createConnection();
			sess = mq.createSession(true, 0);
			prod = sess.createProducer(jms_new_attack);
			
//			BytesMessage jmsMsg = sess.createBytesMessage();
			TextMessage jmsMsg = sess.createTextMessage();
//			jmsMsg.setText(Long.toString(idAttack));
			jmsMsg.setJMSType(AttackInfoType.submithttp.name);
			jmsMsg.setLongProperty("id", attackId);
			jmsMsg.setStringProperty("md5", md5);
			jmsMsg.setStringProperty("sensor_ip", "127.0.0.1"); //TODO
			jmsMsg.setStringProperty("source_ip", form.get(FormFields.source_host.name()));
			jmsMsg.setStringProperty("target_ip", form.get(FormFields.target_host.name()));
			jmsMsg.setStringProperty("md5", md5);
			
			AttackRecord jsonRecord = new AttackRecord();
			jsonRecord.submission_time = "submission_time"; //TODO
			jsonRecord.sensor_ip = "127.0.0.1"; //TODO
			jsonRecord.source_ip = form.get(FormFields.source_host.name());
			jsonRecord.source_port = Integer.parseInt(form.get(FormFields.source_port.name()));
			jsonRecord.target_ip = form.get(FormFields.target_host.name());
			jsonRecord.target_port = Integer.parseInt(form.get(FormFields.target_port.name()));
			jsonRecord.url = form.get(FormFields.url.name());
			jsonRecord.trig = form.get(FormFields.trigger.name());
			jsonRecord.md5 = form.get(FormFields.md5.name());
			jsonRecord.sha512 = form.get(FormFields.sha512.name());
			//TODO remaining fields
			
			ObjectMapper mapper = new ObjectMapper();
//			jmsMsg.writeBytes(mapper.writeValueAsBytes(jsonRecord));
			jmsMsg.setText(mapper.writeValueAsString(jsonRecord));
			
			prod.send(jmsMsg);
			log.debug("scheduled sending '{}' to new_attack", Long.toString(attackId));

//			prod.close();
//			sess.close();
			mq.close();
			db.close();
			
			if (binaryStored) {
				log.info("submission phase 1 complete: {} (stored = {})", SubmitHttpReturnType.S_FILEKNOWN, binaryStored);
				return SubmitHttpReturnType.S_FILEKNOWN;
			}
			else {
				log.info("submission phase 1 complete: {} (stored = {})", SubmitHttpReturnType.S_FILEREQUEST, binaryStored);
				return SubmitHttpReturnType.S_FILEREQUEST;
			}
			
		} catch (SQLException e) {
			log.error("got SQLException", e);
			throw new EJBException(e);
		} catch (JMSException e) {
			log.error("got JMSException", e);
			throw new EJBException(e);
		} catch (JsonGenerationException e) {
			log.error("got JsonGenerationException", e);
			throw new EJBException(e);
		} catch (JsonMappingException e) {
			log.error("got JsonMappingException", e);
			throw new EJBException(e);
		} catch (IOException e) {
			log.error("got IOException", e);
			throw new EJBException(e);
		}
	}
	
	
	private SubmitHttpReturnType handleSecondPhase(HttpServletRequest request) throws SubmissionValidationException {
		log.info("submission phase 2");
		
		String emsg;
		
		Part binaryPart;
		int binarySize;
		String md5;
		
		// get and check md5
		try {
			Part md5Part = request.getPart(FormFields.md5.name());
			long partSize = md5Part.getSize();
			
			if (partSize > MAX_FORM_FIELD_SIZE) {
				emsg = "field too large: " + FormFields.md5.name();
				log.error(emsg);
				throw new SubmissionValidationException(emsg);
			}
			
			byte[] buf = new byte[(int) partSize];
			InputStream is;
			is = md5Part.getInputStream();
			is.read(buf);
			is.close();
			
			md5 = new String(buf);
			
			if (FormFields.md5.validate(md5) == false) {
				emsg = "field does not validate: " + FormFields.md5.name();
				log.error(emsg);
				throw new SubmissionValidationException(emsg);
			}
			
		} catch (IOException e) {
			emsg = "got IOException while trying to read part: " + FormFields.md5.name();
			log.error(emsg, e);
			throw new EJBException(emsg, e);
		} catch (ServletException e) {
			emsg = "got ServletException while trying to read part: " + FormFields.md5.name();
			log.error(emsg, e);
			throw new EJBException(emsg, e);
		} catch (IllegalStateException e) {
			emsg = "got IllegalStateException while trying to read part: " + FormFields.md5.name();
			log.error(emsg, e);
			throw new EJBException(emsg, e);
		}
		
		// get and check binary part
		try {
			binaryPart = request.getPart(BINARY_PART_NAME);
			binarySize = (int) binaryPart.getSize();
			
			if (binarySize == 0) {
				//TODO: S_ERROR
				
			}
			
			if (binaryPart.getContentType().equals(BINARY_CONTENT_TYPE) == false) {
				emsg = "invalid content type for file upload part: " + BINARY_PART_NAME + " - " + binaryPart.getContentType();
				log.error(emsg);
				throw new SubmissionValidationException(emsg);
			}
		} catch (IOException e) {
			emsg = "got IOException while trying read part: " + BINARY_PART_NAME;
			log.error(emsg, e);
			throw new EJBException(emsg, e);
		} catch (ServletException e) {
			emsg = "got ServletException while trying to read part: " + BINARY_PART_NAME;
			log.error(emsg, e);
			throw new EJBException(emsg, e);
		} catch (IllegalStateException e) {
			emsg = "got IllegalStateException while trying to read part: " + BINARY_PART_NAME;
			log.error(emsg, e);
			throw new EJBException(emsg, e);
		}
		
		
		Connection db = null;
		XADiskConnection xadisk = null;
		javax.jms.Connection mq = null;
		Session sess = null;
		MessageProducer prod = null;
		
		PreparedStatement pStmt;
		ResultSet queryRes;
		
		try {
			db = ds.getConnection();
			xadisk = xaDiskConnectionFactory.getConnection();
			
			int binaryId = 0;
			
			pStmt = db.prepareStatement("select binary_id from binaries_pending where md5 = ? limit 1");
			pStmt.setString(1, md5);
			queryRes = pStmt.executeQuery();
			if (queryRes.next()) {
				binaryId = queryRes.getInt(1);
			}
			queryRes.close();
			pStmt.close();
			
			if (binaryId != 0) {
				String path = storageDir;
				for (String pathPart : Storage.pathAsArray(md5)) {
					path += File.separator + pathPart;
					
					File dir = new File(path);
					
					try {
						xadisk.createFile(dir, true);
					} catch (FileAlreadyExistsException e) {
						continue;
					} catch (FileNotExistsException e) {
						// this shouldn't happen
						log.error("parent dir for dir doesn't exist: {}", dir.getAbsolutePath());
						throw new EJBException(e);
					}
				}
				
				// copy the uploaded file to the final storage place (while also checking it's md5)
				File file = new File(path + File.separator + md5 + ".bin");
				try {
					xadisk.createFile(file, false); //TODO set lock timeout?
					XAFileOutputStream out = xadisk.createXAFileOutputStream(file, true);
					InputStream in = binaryPart.getInputStream();
					
					MessageDigest digest = MessageDigest.getInstance("MD5");
					
					
					byte[] buf = new byte[BUF_SIZE];
					
					int n, fileSize = 0;
					while ((n = in.read(buf)) != -1) { //TODO  read timeout?
						log.trace("read loop: read {} bytes", n);
						digest.update(buf, 0, n);
						out.write(buf, 0, n);
						fileSize += n;
					}
					
					out.close();
					String checkedMD5 = new BigInteger(1, digest.digest()).toString(16);
					// pad
					if (checkedMD5.length() < 32) {
						StringBuilder sb = new StringBuilder(32);
						
						int i = 32 - checkedMD5.length();
						while (i > 0) {
							sb.append("0");
							i--;
						}
						
						sb.append(checkedMD5);
						checkedMD5 = sb.toString();
					}
					log.debug("checked md5: {}", checkedMD5);
					
					if (checkedMD5.equals(md5) == false) {
						sctx.setRollbackOnly();
						log.warn("binary integrity check failed");
						return SubmitHttpReturnType.S_ERROR;
					}
					
					pStmt = db.prepareStatement("update binaries set stored = ?, filesize = ? where id = ?");
					pStmt.setBoolean(1, true);
					pStmt.setInt(2, fileSize);
					pStmt.setLong(3, binaryId);
					pStmt.executeUpdate();
					queryRes.close();
					pStmt.close();
					
					pStmt = db.prepareStatement("delete from binaries_pending where binary_id = ?");
					pStmt.setLong(1, binaryId);
					pStmt.executeUpdate();
					queryRes.close();
					pStmt.close();
					
					mq = jmsConnectionFactory.createConnection();
					sess = mq.createSession(true, 0);
					prod = sess.createProducer(jms_new_binary);
					
//					TextMessage jmsMsg = jmsSession.createTextMessage();
					Message jmsMsg = sess.createMessage();
					jmsMsg.setJMSType("unique");
					jmsMsg.setLongProperty("id", binaryId);
					jmsMsg.setStringProperty("md5", md5);
//					jmsMsg.setText(Long.toString(idBinary));
					
					prod.send(jmsMsg);
					log.debug("scheduled sending '{}' to new_binary",	Long.toString(binaryId));

//					prod.close();
//					sess.close();
					mq.close();
					xadisk.close();
					db.close();
					
					log.info("submission phase 2 complete");
					return SubmitHttpReturnType.S_FILEOK;
					
				} catch (FileAlreadyExistsException e) {
					// landing here if a concurrent transaction has copied the
					// same binary just before us
					// or if the other branch could insert a pending entry
					// shortly after they've got deleted in this branch
					// (additional DB/file system consistency checks could be
					// made, but should not be necessary)
					// clean up the pending entries again and return with the
					// appropriate reply (S_FILEOK?)
					pStmt = db.prepareStatement("delete from binaries_pending where binary_id = ?");
					pStmt.setLong(1, binaryId);
					pStmt.executeUpdate();
					queryRes.close();
					pStmt.close();
					
					xadisk.close();
					db.close();
					
					log.debug("binary has been already stored (by a concurrent transaction?) (md5 = {})", md5);
					log.info("submission phase 2 complete");
					return SubmitHttpReturnType.S_FILEOK; // S_FILEKNOWN?
					
				} catch (NoSuchAlgorithmException e) {
					log.error("got NoSuchAlgorithmException", e);
					throw new EJBException(e);
				} catch (IOException e) {
					log.error("got IOException");
					throw new EJBException(e);
				} catch (FileNotExistsException e) {
					// this shouldn't happen
					log.error("parent dir for file doesn't exist: {}", file.getAbsolutePath(), e); // e?
					throw new EJBException(e);
				} catch (FileUnderUseException e) {
					// this shouldn't happen
					log.error("got FileUnderUseException", e);
					throw new EJBException(e);
				} catch (ClosedStreamException e) {
					log.error("got ClosedStreamException", e);
					throw new EJBException(e);
				}
			}
			else {
				db.close();
				xadisk.close();
				log.info("binary is not pending (md5 = {})", md5);
				log.info("submission phase 2 complete");
				return SubmitHttpReturnType.S_FILEOK; // S_FILEKNOWN?
			}
			
		} catch (SQLException e) {
			log.error("got SQLException", e);
			throw new EJBException(e);
		} catch (ResourceException e) {
			log.error("got ResourceException", e);
			throw new EJBException(e);
		} catch (LockingFailedException e) {
			log.error("got LockingFailedException", e);
			throw new EJBException(e);
		} catch (NoTransactionAssociatedException e) {
			log.error("got NoTransactionAssociatedException", e);
			throw new EJBException(e);
		} catch (InsufficientPermissionOnFileException e) {
			log.error("got InsufficientPermissionOnFileException", e);
			throw new EJBException(e);
		} catch (InterruptedException e) {
			log.error("got InterruptedException", e);
			throw new EJBException(e);
		} catch (JMSException e) {
			log.error("got JMSException", e);
			throw new EJBException(e);
		}
		
	}

	
	// converts ipn to string
    private static String ipnToString(String ipn) {
		try {
		    BigInteger ipAsNum = BigInteger.valueOf(new BigInteger(ipn).intValue());
		    InetAddress ip = InetAddress.getByAddress(ipAsNum.toByteArray());
		    return ip.getHostAddress();
		} catch (UnknownHostException e) {
		    log.error("got UnknownHostException while converting ipn", e);
		    throw new EJBException("got UnknownHostException while converting ipn", e);
		}
    }
}
